//
// Created by Sorel on 2021/3/10.
//

#ifndef GA_BASIC_GA_BASIC_HPP
#define GA_BASIC_GA_BASIC_HPP

#include <iostream>
#include <vector>
#include <iomanip>
#include <fstream>

#include "boost_utils.hpp"
#include "cec_utils.hpp"

using namespace std;

namespace GA {

    class Chromosome {  // chromosome in the population
    public:
        vector<double> gene;  // a single gene of the chromosome
        double value;  // value of func will be stored here when evaluated
        double fitness;  // define fitness = 1.0 / (value + 1.0) regarding to minimization problems
    };

    class BasicGA {
    private:
        int pop_size;  // size of the pop
        int dimension;  // dimension of variables
        double cross_prob;
        double mutate_prob;
        double lower_bound;
        double upper_bound;

        vector<Chromosome> population;
        vector<Chromosome> new_population;
        Chromosome pbest;
        Chromosome gbest;
        void print_population(bool show_genes);
        void print_gbest(bool show_genes) const;

        /** basic GA operators **/
        void select();
        vector<double> select_prob;
        void cross();
        void single_point_cross(int p1, int p2);
        void mutate();
        void elitist();
        void evaluate(int func_num);

        string save_path;
        vector<double> value_list;  // record the history of gbest fitness
        void save_results();

    public:
        void set_save_path(const string& _save_path);
        void init(int _pop_size, int _dimension, double _cross_prob, double _mutate_prob);
        void set_bounds(double _lower, double _upper);
        void execute(int func_num, int depth, int trial);
    };

    void BasicGA::init(int _pop_size, int _dimension, double _cross_prob, double _mutate_prob) {
        this->pop_size = _pop_size;
        this->dimension = _dimension;
        this->cross_prob = _cross_prob;
        this->mutate_prob = _mutate_prob;
        population.resize(pop_size);
        new_population.resize(pop_size);
        for (int p = 0; p < pop_size; ++p) {
            population[p].gene.resize(dimension + 1);
            new_population[p].gene.resize(dimension + 1);
        }
        pbest.gene.resize(dimension + 1);
        gbest.gene.resize(dimension + 1);
    }

    void BasicGA::set_bounds(double _lower, double _upper) {
        this->lower_bound = _lower;
        this->upper_bound = _upper;
    }

    void BasicGA::execute(int func_num, int depth, int trial) {
        value_list.clear();
        value_list.resize(depth);
        fill(value_list.begin(), value_list.end(), 0);
        /** the main GA procedure starts **/
        int _trial = 0;
        cout<<"trial ";
        while (_trial < trial) {
            // init population
            for (int p = 0; p < pop_size; ++p) {
                for (int i = 0; i < dimension; ++i) {
                    population[p].gene[i] = random_double(lower_bound, upper_bound);  // random initialization
                }
            }
            gbest.fitness = pbest.fitness = 0.0;
            evaluate(func_num);  // evaluate population
            int gen = 0;
            while (gen < depth) {
                select();
                cross();
                mutate();
                elitist();
                evaluate(func_num);
                value_list[gen] += gbest.value;
                gen++;
            }
            _trial++;
            cout<<_trial<<", ";
        } cout<<" all trials completed. ";
        /** the main GA procedure ends **/
        for (auto &value : value_list) {
            value = value / trial;
        }
        save_results();
    }

    void BasicGA::select() {
        // select by fitness proportional model
        int p;
        select_prob.clear();
        double sum = 0;
        for (p = 0; p < pop_size; ++p) {
            sum += population[p].fitness;
        }
        for (p = 0; p < pop_size; ++p) {
            select_prob.push_back(population[p].fitness/sum);
        }
        for (p = 1; p < pop_size; ++p) {
            select_prob[p] = select_prob[p] + select_prob[p-1];
        }
        double prob;
        new_population.clear();
        new_population.resize(pop_size);
        for (int i = 0; i < pop_size; ++i) {
            prob = random_p();
            for (p = 0; p < pop_size; ++p) {
                if (prob <= select_prob[p]) {
                    new_population[i].gene.assign(population[p].gene.begin(), population[p].gene.end());
                    new_population[i].fitness = population[p].fitness;
                    break;
                }
            }
        }
        for (p = 0; p < pop_size; ++p) {
            population.assign(new_population.begin(), new_population.end());
        }
    }

    void BasicGA::cross() {
        int p, p1, p2;
        int pool = 0;  // mating pool
        double prob;
        for (p = 0; p < pop_size; ++p)
        {
            prob = random_p();
            if (prob < cross_prob) {
                pool++;
                p1 = p;
                if (pool % 2 == 0)
                    single_point_cross(p1, p2);
                else
                    p2 = p1;
            }
        }
    }

    void BasicGA::single_point_cross(int p1, int p2) {
        int point = random_int(0, dimension);
        for (int i = point+1; i < dimension; ++i) {
            swap(population[p1].gene[i], population[p2].gene[i]);  // swap genes from point+1 to the end
        }
    }

    void BasicGA::mutate() {
        double prob;
        int i;
        for (int p = 0; p < pop_size; ++p) {
            for (i = 0; i < dimension; ++i) {
                prob = random_p();
                if (prob <= mutate_prob) {
                    population[p].gene[i] = random_double(lower_bound, upper_bound);
                }
            }
        }
    }

    void BasicGA::elitist() {
        // replace the worst individual with pbest from last generation, performing before evaluation
        int flag = 0;
        double fitness = population[0].fitness;
        for (int p = 0; p < pop_size; ++p) {
            if (population[p].fitness < fitness) {
                fitness = population[p].fitness;
                flag = p;
            }
        }
        population[flag].gene.assign(pbest.gene.begin(), pbest.gene.end());
        population[flag].fitness = pbest.fitness;
    }

    void BasicGA::evaluate(int func_num) {
        for (int p = 0; p < pop_size; ++p) {
            /** replace your evaluation operator here **/
            cec14_test_func(population[p].gene.data(), &population[p].value, dimension, 1, func_num);

            population[p].fitness = 1.0 / (population[p].value + 1.0);  // minimization problem
            if (pbest.fitness < population[p].fitness) {  // update pbest
                pbest.fitness = population[p].fitness;
                pbest.value = population[p].value;
                pbest.gene.assign(population[p].gene.begin(), population[p].gene.end());
            }
        }
        if (gbest.fitness < pbest.fitness) {  // update gbest
            gbest.fitness = pbest.fitness;
            gbest.value = pbest.value;
            gbest.gene.assign(pbest.gene.begin(), pbest.gene.end());
        }
    }

    void BasicGA::set_save_path(const string& _save_path) {
        this->save_path = _save_path;
        string file_name = save_path;
        ofstream out_file(file_name);
        if (!out_file.is_open()) {
            cerr<<"save_path invalid, please check directories. "<<endl;
            exit(-1);
        }
        out_file.close();
    }

    void BasicGA::save_results() {
        string file_name = save_path;
        ofstream out_file(file_name);
        if (!out_file.is_open()) {
            cerr<<"save_path invalid, please check directories. "<<endl;
            exit(-1);
        }
        for (double value : value_list) {
            out_file<<value<<endl;
        }
        out_file.close();
        cout<<save_path<<" created. "<<endl;
    }

    void BasicGA::print_population(bool show_genes) {
        cout<<"===================== population ====================="<<endl;
        int count = 0;
        for (auto p : population) {
            cout<<"p "<<count+1<<": ";
            if (show_genes) {
                for (int i = 0; i < dimension; ++i) {
                    cout<<p.gene[i]<<", ";
                }
            } cout<<"fitness: "<<p.fitness<<". "<<endl;
        }
    }

    void BasicGA::print_gbest(bool show_genes) const {
        cout<<"gbest: ";
        if (show_genes) {
            for (int i = 0; i < dimension; ++i) {
                cout<<gbest.gene[i]<<", ";
            }
        } cout<<"fitness: "<<gbest.fitness<<". "<<endl;
    }
}

#endif //GA_BASIC_GA_BASIC_HPP
